"""
QUADRATIC utilities
"""
import array
import numpy as np
import tensorflow.compat.v1 as tf
#tf.compat.v1.disable_eager_execution()

class Data(object):
    """
    Data class takes care of the minibatch management

    There are two optional parameters for Data objects.
    -- batch_size: if lefts as None will use the entire data set for 
       each pass.
    -- shuffle: wich defaults to true will shuffle data points to feed 
       the network at each iteration. When a epoch has been reached, 
       data is reshuffled.
    """
    batch_iter = None 
    data_size = None
    batch_idx = None
    crt_idx = None
    n_batches = None
    shuffle = None
    
    def __init__(self, data, batch_size=None, shuffle=True):
        self.data = data
        self.data_size = data.shape[0]
        self.shuffle = shuffle
        if batch_size is None:
            self.batch_size = self.data_size
            self.batch_iter = None
            self.batch_idx = None
            self.shuffle = False
        else:
            err_msg = "'batch_size exeeds number of data samples'"
            assert self.data_size >= batch_size, err_msg
            self.batch_size = batch_size
            self.batch_iter = 0
        self.n_batches = np.int64(self.data_size / self.batch_size)
        leftovers = self.data_size % self.batch_size
                
    def getBatch(self, batch_size):
        if self.batch_iter is None:
            return self.data
  
        if self.batch_iter == 0:
            ## shuffle data at the begining of every epoch
            if self.shuffle is True:
                self.batch_idx = np.random.permutation(self.data_size)
            else:
                self.batch_idx = np.arange(self.data_size, dtype=np.int64)
            self.batch_idx = np.reshape(self.batch_idx, (self.n_batches, -1))
        self.crt_idx = self.batch_idx[self.batch_iter, :]
        batch = self.data[self.crt_idx, ::]
        self.batch_iter += 1
        if self.batch_iter == self.n_batches:
            self.batch_iter = 0
        return batch
    

def ZCAMatrix(cov_matrix, zeps=1e-10):
    """
    Construct Zero phase componenet analysis transformation for whitening data
    
    Inputs:
    -- cov_matrix : 2-dimensional numpy array containign the covariance matrix 
    from which the whitening transformation will be constructed.
    -- zeps : (optional) scalar to regularize the inverse of the eigenvalues of 
    the covariance matrix employed in the ZCA transform.
    
    Outputs:
    -- Q : numpy array that transforms data points expressed as columns into 
    their whitened versions.
    """
    eigval, eigvec = np.linalg.eig(cov_matrix)
    Q = np.dot(eigvec/(np.sqrt(eigval) + zeps), eigvec.transpose())
    return Q

